#!/usr/bin/env python

from __future__ import print_function

from sensor_msgs.msg import JointState
from geometry_msgs.msg import Twist
from nav_msgs.msg import Odometry

import rospy
import actionlib
from control_msgs.msg import FollowJointTrajectoryAction
from control_msgs.msg import FollowJointTrajectoryGoal
from trajectory_msgs.msg import JointTrajectoryPoint

from sensor_msgs.msg import PointCloud2

from std_srvs.srv import Trigger, TriggerRequest, TriggerResponse

import math
import time
import threading
import sys

import tf2_ros
import argparse as ap

import hello_helpers.hello_misc as hm
import stretch_funmap.navigate as nv

class OpenDrawerNode(hm.HelloNode):

    def __init__(self):
        hm.HelloNode.__init__(self)
        self.rate = 10.0
        self.joint_states = None
        self.joint_states_lock = threading.Lock()
        self.move_base = nv.MoveBase(self)
        self.letter_height_m = 0.2
        self.wrist_position = None
        self.lift_position = None
        
    def joint_states_callback(self, joint_states):
        with self.joint_states_lock: 
            self.joint_states = joint_states
        wrist_position, wrist_velocity, wrist_effort = hm.get_wrist_state(joint_states)
        self.wrist_position = wrist_position        
        lift_position, lift_velocity, lift_effort = hm.get_lift_state(joint_states)
        self.lift_position = lift_position
        
    def align_to_surface(self):
        rospy.loginfo('align_to_surface')
        trigger_request = TriggerRequest() 
        trigger_result = self.trigger_align_with_nearest_cliff_service(trigger_request)
        rospy.loginfo('trigger_result = {0}'.format(trigger_result))

    def extend_hook_until_contact(self):
        rospy.loginfo('extend_hook_until_contact')
        max_extension_m = 0.5
        max_reach_m = 0.4
        extension_m = self.wrist_position + max_reach_m
        extension_m = min(extension_m, max_extension_m)
        extension_contact_effort = 45.0 #42.0 #40.0 from funmap
        pose = {'wrist_extension': (extension_m, extension_contact_effort)}
        self.move_to_pose(pose, custom_contact_thresholds=True)
        
    def lower_hook_until_contact(self):
        rospy.loginfo('lower_hook_until_contact')
        max_drop_m = 0.15
        lift_m = self.lift_position - max_drop_m
        lift_contact_effort = 16.0 #18.0 #20.0 #20.0 from funmap
        pose = {'joint_lift': (lift_m, lift_contact_effort)}
        self.move_to_pose(pose, custom_contact_thresholds=True)

        use_correction = True
        if use_correction:
            # raise due to drop down after contact detection
            #rospy.sleep(0.5) # wait for new lift position
            rospy.sleep(0.2) # wait for new lift position
            lift_m = self.lift_position + 0.015
            pose = {'joint_lift': lift_m}
            self.move_to_pose(pose)
            rospy.sleep(0.2) # wait for new lift position

    def raise_hook_until_contact(self):
        rospy.loginfo('raise_hook_until_contact')
        max_raise_m = 0.15
        lift_m = self.lift_position + max_raise_m
        lift_contact_effort = 45.0
        pose = {'joint_lift': (lift_m, lift_contact_effort)}
        self.move_to_pose(pose, custom_contact_thresholds=True)

        #rospy.sleep(1.0) # needed to correct for bounce after contact
        
        use_correction = True
        if use_correction:
            # raise due to drop down after contact detection
            rospy.sleep(0.5) # wait for new lift position
            lift_m = self.lift_position + 0.01 #0.015
            pose = {'joint_lift': lift_m}
            self.move_to_pose(pose)
            rospy.sleep(0.5) # wait for new lift position
        
    def backoff_from_surface(self):
        rospy.loginfo('backoff_from_surface')
        if self.wrist_position is not None:
            wrist_target_m = self.wrist_position - 0.005
            pose = {'wrist_extension': wrist_target_m}
            self.move_to_pose(pose)
            return True
        else:
            rospy.logerr('backoff_from_surface: self.wrist_position is None!')
            return False

    def pull_open(self):
        rospy.loginfo('pull_open')
        if self.wrist_position is not None:
            max_extension_m = 0.5
            extension_m = self.wrist_position - 0.2
            extension_m = min(extension_m, max_extension_m)
            extension_m = max(0.01, extension_m)
            extension_contact_effort = 120 #100.0 #40.0 from funmap
            pose = {'wrist_extension': (extension_m, extension_contact_effort)}
            self.move_to_pose(pose, custom_contact_thresholds=True)
            return True
        else:
            rospy.logerr('pull_open: self.wrist_position is None!')
            return False

    def push_closed(self):
        rospy.loginfo('push_closed')
        if self.wrist_position is not None:
            wrist_target_m = self.wrist_position + 0.22
            pose = {'wrist_extension': wrist_target_m}
            self.move_to_pose(pose)
            return True
        else:
            rospy.logerr('pull_open: self.wrist_position is None!')
            return False
        
    def move_to_initial_configuration(self):
        initial_pose = {'wrist_extension': 0.01,
                        'joint_wrist_yaw': 1.570796327,
                        'gripper_aperture': 0.0}
        rospy.loginfo('Move to the initial configuration for drawer opening.')
        self.move_to_pose(initial_pose)

    def trigger_open_drawer_down_callback(self, request):
        return self.open_drawer('down')
        
    def trigger_open_drawer_up_callback(self, request):
        return self.open_drawer('up')


    def open_drawer(self, direction):
        self.move_to_initial_configuration()

        #self.align_to_surface()
        self.extend_hook_until_contact()
        success = self.backoff_from_surface()
        if not success:
            return TriggerResponse(
                success=False,
                message='Failed to backoff from the surface.'
            )

        if direction == 'down': 
            self.lower_hook_until_contact()
        elif direction == 'up':
            self.raise_hook_until_contact()
            
        success = self.pull_open()
        if not success:
            return TriggerResponse(
                success=False,
                message='Failed to pull open the drawer.'
            )

        push_drawer_closed = False
        if push_drawer_closed: 
            rospy.sleep(3.0)
            self.push_closed()
        
        return TriggerResponse(
            success=True,
            message='Completed opening the drawer!'
            )

    
    def main(self):
        hm.HelloNode.main(self, 'open_drawer', 'open_drawer', wait_for_first_pointcloud=False)

        self.joint_states_subscriber = rospy.Subscriber('/stretch/joint_states', JointState, self.joint_states_callback)
        
        self.trigger_open_drawer_service = rospy.Service('/open_drawer/trigger_open_drawer_down',
                                                         Trigger,
                                                         self.trigger_open_drawer_down_callback)

        self.trigger_open_drawer_service = rospy.Service('/open_drawer/trigger_open_drawer_up',
                                                         Trigger,
                                                         self.trigger_open_drawer_up_callback)
        
        
        rospy.wait_for_service('/funmap/trigger_align_with_nearest_cliff')
        rospy.loginfo('Node ' + self.node_name + ' connected to /funmap/trigger_align_with_nearest_cliff.')
        self.trigger_align_with_nearest_cliff_service = rospy.ServiceProxy('/funmap/trigger_align_with_nearest_cliff', Trigger)

        rospy.wait_for_service('/funmap/trigger_reach_until_contact')
        rospy.loginfo('Node ' + self.node_name + ' connected to /funmap/trigger_reach_until_contact.')
        self.trigger_reach_until_contact_service = rospy.ServiceProxy('/funmap/trigger_reach_until_contact', Trigger)

        rospy.wait_for_service('/funmap/trigger_lower_until_contact')
        rospy.loginfo('Node ' + self.node_name + ' connected to /funmap/trigger_lower_until_contact.')
        self.trigger_lower_until_contact_service = rospy.ServiceProxy('/funmap/trigger_lower_until_contact', Trigger)
        
        rate = rospy.Rate(self.rate)
        while not rospy.is_shutdown():
            rate.sleep()

        
if __name__ == '__main__':
    try:
        parser = ap.ArgumentParser(description='Open Drawer behavior for stretch.')
        #parser.add_argument('--mapping_on', action='store_true', help='Turn on mapping control. For example, the space bar will trigger a head scan. This requires that the mapping node be run (funmap).')
        args, unknown = parser.parse_known_args()
        node = OpenDrawerNode()
        node.main()
    except KeyboardInterrupt:
        rospy.loginfo('interrupt received, so shutting down')
#    except rospy.ROSInterruptException:
#        pass
        #rospy.loginfo('keyboard_teleop was interrupted', file=sys.stderr)

